package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"path"
	"runtime"
	"strings"
	"time"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
	// "golang.org/x/crypto/ssh/terminal"
)

/****************************************
windows下编译成linux的方法:
set CGO_ENABLED=0
set GOOS=linux
set GOARCH=amd64
go build -o issh main.go

参数说明:
# -H 是ssh的远程地址，默认是localhost
# -P 是端口号，默认是22
# -u 是用户，默认root
# -p 是密码，帮助中不打印，但默认是xxxxxx
# -c 是要执行的命令
# -t 启动Terminal后执行的命令，具备交互功能
# -f 是文件上传下载功能

使用方法：
issh -H 123.124.125.126 -P 22 -u root -p xxxxxx -c ls  # 执行ls命令，执行命令后直接返回结果
issh -H 123.124.125.126 -P 22 -u root -p xxxxxx -t bash # 相当于ssh登录到远程

issh -H 123.124.125.126 -P 22 -u root -p xxxxxx -f "upload:./nohup.out|/root"
issh -H 123.124.125.126 -P 22 -u root -p xxxxxx -f "download:/root/nohup.out|/root"
# -f是文件操作，
# upload是上传，第一个路径是本地源文件(或文件夹，文件夹时会遍历文件夹下的所有文件和子文件夹)路径，第二个路径是远程目标文件夹路径(如果目标文件夹没有会自动创建)，
# download是下载，第一个路径是远程源文件(或文件夹，文件夹时会遍历文件夹下的所有文件和子文件夹)路径，第二个路径是本地目标文件夹路径(如果目标文件夹没有会自动创建)
# 本地路径可以使用相对路径，远程路径也可以使用相对用户的路径

****************************************/
type sliceFlag []string

func (f *sliceFlag) String() string {
	return fmt.Sprintf("%v", []string(*f))
}

func (f *sliceFlag) Set(value string) error {
	*f = append(*f, value)
	return nil
}

type Cli struct {
	IP         string
	Username   string
	Password   string
	Port       int
	sshclient  *ssh.Client  //ssh Client
	ftpclient  *sftp.Client //sftp Client
	LastResult string       //last Run Result
}

func (c *Cli) Run(shell string) (string, error) {
	if c.sshclient == nil {
		if err := c.connect(); err != nil {
			return "", err
		}
	}
	session, err := c.sshclient.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()
	buf, err := session.CombinedOutput(shell)

	c.LastResult = string(buf)
	return c.LastResult, err
}
func (c *Cli) connect() error {
	config := ssh.ClientConfig{
		User: c.Username,
		Auth: []ssh.AuthMethod{ssh.Password(c.Password)},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
		Timeout: 10 * time.Second,
	}
	addr := fmt.Sprintf("%s:%d", c.IP, c.Port)
	sshClient, err := ssh.Dial("tcp", addr, &config)
	if err != nil {
		return err
	}
	c.sshclient = sshClient
	return nil
}
func (c *Cli) close() error {
	if c.ftpclient != nil {
		err := c.ftpclient.Close()
		if err != nil {
			return err
		}
	}
	if c.sshclient != nil {
		err := c.sshclient.Close()
		if err != nil {
			return err
		}
	}
	return nil
}
func (c *Cli) RunTerminal(shell string, stdin io.Reader, stdout, stderr io.Writer) error {
	if runtime.GOOS == "windows" {
		return errors.New("windows does not support terminal")
	}
	if c.sshclient == nil {
		if err := c.connect(); err != nil {
			return err
		}
	}
	session, err := c.sshclient.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	fd := int(os.Stdin.Fd())
	// oldState, err := terminal.MakeRaw(fd)
	oldState, err := term.MakeRaw(fd)
	if err != nil {
		panic(err)
	}
	// defer terminal.Restore(fd, oldState)
	defer term.Restore(fd, oldState)

	session.Stdout = stdout
	session.Stderr = stderr
	session.Stdin = stdin

	// termWidth, termHeight, err := terminal.GetSize(fd)
	termWidth, termHeight, err := term.GetSize(fd)
	if err != nil {
		panic(err)
	}
	// Set up terminal modes
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,     // enable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	// Request pseudo terminal
	if err := session.RequestPty("xterm-256color", termHeight, termWidth, modes); err != nil {
		return err
	}

	session.Run(shell)
	return nil
}

type WriteCounter struct {
	startTime  int64
	endTime    int64
	FileSize   int64
	WriteTotal int64
}

func (wc *WriteCounter) Write(p []byte) (int, error) {
	n := len(p)
	wc.WriteTotal += int64(n)
	tmpEndTime := time.Now().UnixMicro()
	speed := float64(n) / 1024.0 * 1e6 / float64(tmpEndTime-wc.endTime)
	fmt.Printf("\r  %.f%%   %dB/%dB  %.2fKB/s  %s  %s     ",
		float64(wc.WriteTotal)/float64(wc.FileSize)*100.0,
		wc.WriteTotal,
		wc.FileSize,
		speed,
		time.UnixMicro(tmpEndTime-wc.startTime).UTC().Format("15:04:05"),
		time.UnixMicro(int64(float64(wc.FileSize-wc.WriteTotal)/speed*1024)).UTC().Format("15:04:05"))
	wc.endTime = tmpEndTime
	return n, nil
}

/* UploadFile 将本地文件上传到远程服务器
  localfile: 本地文件或者文件夹路径
  remotePath: 远程文件夹路径
  identify: 上传模式
       0:直接覆盖
	   1:如果文件大小不同或者修改时间较新则覆盖，否则略过
	   2:如果远程文件已经存在,则询问是否覆盖
*/
func (c *Cli) UploadFile(localfile string, remotePath string, identify int) error {
	//log.Printf("Files upload:%v--->%v ---start\n", localfile, remotePath)
	localFileStat, err := os.Stat(localfile)
	if err != nil {
		errmsg := fmt.Sprintf("Local file path %v, error:%v\n", localfile, err)
		return errors.New(errmsg)
	}
	if c.sshclient == nil {
		if err := c.connect(); err != nil {
			return err
		}
	}
	if c.ftpclient == nil {
		ftpclient, err := sftp.NewClient(c.sshclient)
		if err != nil {
			log.Println("sftp client create error,", err)
			return err
		}
		c.ftpclient = ftpclient
	}

	if localFileStat.IsDir() { // 如果本地是文件夹
		remoteDirStat, err := c.ftpclient.Stat(remotePath)
		if err != nil { // 如果远程文件路径不存在则创建文件夹路径
			rerr := c.ftpclient.MkdirAll(remotePath)
			if rerr != nil {
				errmsg := fmt.Sprintf("Create remote path %v fail,error:%v\n", remotePath, rerr)
				return errors.New(errmsg)
			}
			rerr = c.UploadFile(localfile, remotePath, identify)
			return rerr
		}
		if remoteDirStat.IsDir() { // 如果远程也是个文件夹，则遍历本地文件夹上传文件
			localFiles, err := ioutil.ReadDir(localfile)
			if err != nil {
				log.Fatalf("read local path %v filelist fail,error:%v\n", localfile, err)
			}
			for _, backupDir := range localFiles {
				localFilePath := path.Join(localfile, backupDir.Name())
				remoteFilePath := path.Join(remotePath, backupDir.Name())
				if backupDir.IsDir() {
					merr := c.UploadFile(localFilePath, remoteFilePath, identify)
					if merr != nil {
						return merr
					}
				} else {
					merr := c.UploadFile(localFilePath, remotePath, identify)
					if merr != nil {
						return merr
					}
				}
			}
		} else { // 如果本地是个文件夹，远程是个文件，则不能上传
			errmsg := fmt.Sprintf("Create remote path %v fail,error:remote path %v is a file, can not create remove Directory\n", localfile, remotePath)
			return errors.New(errmsg)
		}
	} else { // 如果本地是文件
		remoteDirStat, err := c.ftpclient.Stat(remotePath)
		if err != nil { // 如果远程文件夹不存在，则创建远程文件夹
			err := c.ftpclient.MkdirAll(remotePath)
			if err != nil {
				log.Printf("Create Remote Dir %v error:%v\n", remotePath, err)
				return err
			}
			c.UploadFile(localfile, remotePath, identify)
		} else {
			if remoteDirStat.IsDir() { // 如果本地是文件,远程是文件夹,则上传文件夹
				// log.Println("remote Directory", remotePath, " exist, not create dir")
				remoteFileName := path.Base(localfile)
				remotefile := path.Join(remotePath, remoteFileName)

				remoteFileStat, err := c.ftpclient.Stat(remotefile)
				if err == nil { // 远程已经存在了要保存成的文件或者文件夹了，如果是文件夹则不能上传，如果是文件则将远程文件备份后上传
					if remoteFileStat.IsDir() {
						errmsg := fmt.Sprintf("remote %v is directory,Can not create file", remotefile)
						return errors.New(errmsg)
					} else {
						if identify == 0 { // 覆盖远程文件
							errmsg := fmt.Sprintf("Remote file %v already exists, the old file will be rewrite", remotefile)
							log.Println(errmsg)
						} else if identify == 1 { // 如果文件大小不同或者修改时间较新则覆盖，否则略过
							if localFileStat.Size() != remoteFileStat.Size() || remoteDirStat.ModTime().Local().Before(localFileStat.ModTime().Local()) {
								errmsg := fmt.Sprintf("Remote file %v will be rewrite", remotefile)
								log.Println(errmsg)
							} else {
								errmsg := fmt.Sprintf("Remote file %v exists, Do not upload", remotefile)
								log.Println(errmsg)
								return nil
							}
						} else if identify == 2 {
							var userinput string
							fmt.Printf("local:%s --> remote:%s (Y/N)\n", localfile, remotefile)
							for {
								fmt.Scanln(&userinput)
								if userinput == "Y" || userinput == "y" || userinput == "N" || userinput == "n" {
									break
								}
							}
							if userinput == "Y" || userinput == "y" {
								return nil
							}
						} else {
							errmsg := fmt.Sprintf("identify(%d) is not a valid parameter", identify)
							return errors.New(errmsg)
						}
					}
				}

				// 正式上传文件
				srcFile, err := os.Open(localfile)
				if err != nil {
					log.Println("open local file fail,", err)
					// panic(err)
					return err
				}
				defer srcFile.Close()
				dstFile, e := c.ftpclient.Create(remotefile)
				if e != nil {
					log.Println("create remote file fail", err)
					return e
				}
				defer dstFile.Close()

				// _, err = dstFile.ReadFrom(srcFile)
				// if err != nil {
				// if _, err := io.Copy(dstFile, srcFile); err != nil {
				// 	log.Println("read local file fail", err)
				// 	return err
				// }

				fmt.Printf("Uploading %s to %s\n", localfile, remotefile)
				sttime := time.Now().UnixMicro()
				counter := &WriteCounter{FileSize: localFileStat.Size(), startTime: sttime, endTime: sttime}
				upFileSize, err := io.Copy(dstFile, io.TeeReader(srcFile, counter))
				if err != nil {
					log.Println("read local file fail", err)
					return err
				}

				// 修改远程文件的权限
				dstFile.Chmod(localFileStat.Mode())
				// 修改远程文件的时间
				c.ftpclient.Chtimes(remotefile, localFileStat.ModTime(), localFileStat.ModTime())

				fmt.Println()
				fmt.Printf("%s: %d bytes transferred in %d seconds (%.f KB/s)\n", localfile, upFileSize, (counter.endTime-counter.startTime)/1e6, float64(upFileSize)/1024.0*1e6/float64(counter.endTime-counter.startTime))
				// fmt.Printf("----上传文件:本地文件大小%v,上传文件大小%v", localFileStat.Size(), write_size)
				// buffer := make([]byte, 1024)
				// for {
				// 	n, err := srcFile.Read(buffer)
				// 	if err != nil {
				// 		if err == io.EOF {
				// 			// log.Println("已读取到文件末尾")
				// 			break
				// 		} else {
				// 			log.Println("read local file fail", err)
				// 			return err
				// 		}
				// 	}
				// 	dstFile.Write(buffer[:n])
				// 	//注意，由于文件大小不定，不可直接使用buffer，否则会在文件末尾重复写入，以填充1024的整数倍
				// }
			} else { // 如果本地是个文件，远程是个文件，不能上传
				// log.Fatalf("Remote %v is file,Can not create file", remotePath)
				errmsg := fmt.Sprintf("Remote file %v exist,rewrite remote file", remotePath)
				return errors.New(errmsg)
			}
		}
	}
	//log.Printf("Files upload:%v--->%v ---sucessfull\n", localfile, remotePath)
	return nil
}

/* DownloadFile 将远程服务器上的文件或文件夹下载到本地
  localfile: 本地文件夹路径
  remotePath: 远程文件或文件夹路径
  identify: 下载模式
       0:直接覆盖
	   1:如果文件大小不同或者修改时间较新则覆盖，否则略过
	   2:如果远程文件已经存在,则询问是否覆盖
*/
func (c *Cli) DownloadFile(localPath string, remoteFile string, identify int) error {
	//log.Printf("Files download:%v--->%v ---start\n", remoteFile, localPath)
	if c.sshclient == nil {
		if err := c.connect(); err != nil {
			return err
		}
	}
	if c.ftpclient == nil {
		ftpclient, err := sftp.NewClient(c.sshclient)
		if err != nil {
			log.Println("sftp client create error,", err)
			panic(err)
		}
		c.ftpclient = ftpclient
	}

	remoteFileStat, err := c.ftpclient.Stat(remoteFile)
	if err != nil {
		errmsg := fmt.Sprintf("Remote file path %v, error:%v\n", remoteFile, err)
		return errors.New(errmsg)
	}
	if remoteFileStat.IsDir() {
		// 远程文件是个文件夹
		localDirStat, err := os.Stat(localPath)
		if err != nil {
			// 本地找不到路径，创建路径，要判断一下错误是不是文件不存在
			merr := os.MkdirAll(localPath, os.ModePerm)
			if merr != nil {
				errmsg := fmt.Sprintf("Failed to create local path %v", localPath)
				return errors.New(errmsg)
			}
			merr = c.DownloadFile(localPath, remoteFile, identify)
			return merr
		} else {
			if localDirStat.IsDir() {
				// 远程是个文件夹,本地路径是个文件夹,遍历下载
				listRemoteFile, err := c.ftpclient.ReadDir(remoteFile)
				if err != nil {
					errmsg := fmt.Sprintf("Read remote Director %v fail, %v", remoteFile, err)
					return errors.New(errmsg)
				}
				for _, backupDir := range listRemoteFile {
					localFilePath := path.Join(localPath, backupDir.Name())
					remoteFilePath := path.Join(remoteFile, backupDir.Name())
					if backupDir.IsDir() {
						merr := c.DownloadFile(localFilePath, remoteFilePath, identify)
						if merr != nil {
							return merr
						}
					} else {
						merr := c.DownloadFile(localPath, remoteFilePath, identify)
						if merr != nil {
							return merr
						}
					}
				}

			} else {
				// 远程是个文件夹，本地路径是个文件,不能下载
				errmsg := fmt.Sprintf("Local path %v is file, cannt download directory %v", localPath, remoteFile)
				return errors.New(errmsg)
			}
		}
	} else { // 远程文件是个文件
		localDirStat, err := os.Stat(localPath)
		if err != nil { // 本地找不到路径，创建路径
			merr := os.MkdirAll(localPath, os.ModePerm)
			if merr != nil {
				errmsg := fmt.Sprintf("Failed to create local path %v", localPath)
				return errors.New(errmsg)
			}
			merr = c.DownloadFile(localPath, remoteFile, identify)
			return merr
		} else {
			if localDirStat.IsDir() { // 远程是个文件,本地路径是个文件夹,可以下载
				remoteFilename := path.Base(remoteFile)
				localfile := path.Join(localPath, remoteFilename)
				localFileStat, merr := os.Stat(localfile)
				if merr == nil {
					if localFileStat.IsDir() {
						errmsg := fmt.Sprintf("Local path %v is directory, cannt download %v", remoteFile, localfile)
						return errors.New(errmsg)
					} else {
						if identify == 0 { // 覆盖远程文件
							errmsg := fmt.Sprintf("Local file %v already exists, the old file will be rewrite", localfile)
							log.Println(errmsg)
						} else if identify == 1 { // 如果文件大小不同或者修改时间较新则覆盖，否则略过
							if localFileStat.Size() != remoteFileStat.Size() || localFileStat.ModTime().Local().Before(remoteFileStat.ModTime().Local()) {
								errmsg := fmt.Sprintf("Local file %v will be rewrite", localfile)
								log.Println(errmsg)
							} else {
								errmsg := fmt.Sprintf("Remote file %v exists, Do not upload", localfile)
								log.Println(errmsg)
								return nil
							}
						} else if identify == 2 {
							var userinput string
							fmt.Printf("remote:%s --> local:%s (Y/N)\n", remoteFile, localfile)
							for {
								fmt.Scanln(&userinput)
								if userinput == "Y" || userinput == "y" || userinput == "N" || userinput == "n" {
									break
								}
							}
							if userinput == "Y" || userinput == "y" {
								return nil
							}
						} else {
							errmsg := fmt.Sprintf("identify(%d) is not a valid parameter", identify)
							return errors.New(errmsg)
						}

					}
				}

				// 正式下载文件
				srcFile, err := c.ftpclient.Open(remoteFile)
				if err != nil {
					errmsg := fmt.Sprintf("read remote file %v fail", err)
					return errors.New(errmsg)
				}
				defer srcFile.Close()
				dstFile, e := os.Create(localfile)
				if e != nil {
					errmsg := fmt.Sprintf("create local file %v fail,%v", localfile, e)
					return errors.New(errmsg)
				}
				defer dstFile.Close()
				// if _, err1 := srcFile.WriteTo(dstFile); err1 != nil {
				// if _, err := io.Copy(dstFile, srcFile); err != nil {
				// 	errmsg := fmt.Sprintf("local file %v write fail,", err)
				// 	return errors.New(errmsg)
				// }
				fmt.Printf("Downloading %s to %s\n", remoteFile, localfile)
				sttime := time.Now().UnixMicro()
				counter := &WriteCounter{FileSize: remoteFileStat.Size(), startTime: sttime, endTime: sttime}
				downFileSize, err := io.Copy(dstFile, io.TeeReader(srcFile, counter))
				if err != nil {
					log.Println("read local file fail", err)
					return err
				}
				// 修改本地文件权限
				dstFile.Chmod(remoteFileStat.Mode())
				// 修改本地文件时间
				os.Chtimes(localfile, remoteFileStat.ModTime(), remoteFileStat.ModTime())
				fmt.Println()
				fmt.Printf("%s: %d bytes transferred in %d seconds (%.f KB/s)\n", remoteFile, downFileSize, (counter.endTime-counter.startTime)/1e6, float64(downFileSize)/1024.0*1e6/float64(counter.endTime-counter.startTime))
			} else { // 本地路径是个文件,不能下载，报错
				errmsg := fmt.Sprintf("Local path %v is file, rewrite local file", localPath)
				return errors.New(errmsg)
			}
		}
	}
	//log.Printf("Files download:%v--->%v ---sucessfull\n", remoteFile, localPath)
	return nil
}

func main() {
	log.SetFlags(log.Ldate | log.Ltime)
	var host string
	var port int
	var user string
	var password string
	var TerminalCmd string
	var cmdList sliceFlag
	var fileoptlist sliceFlag
	var identify int
	flag.StringVar(&host, "H", "localhost", "ssh host")
	flag.IntVar(&port, "P", 22, "ssh port")
	flag.StringVar(&user, "u", "root", "ssh user")
	flag.StringVar(&password, "p", "", "ssh password")
	flag.Var(&cmdList, "c", "Application hosts,for example: -c=\"cd /root;pwd\"")
	flag.StringVar(&TerminalCmd, "t", "", "start Terminal and run this cmd,example:bash")
	flag.Var(&fileoptlist, "f", `file upload,args[upload:localfile:remotepath],example: upload:/root/a.txt|/root
file download,args[download:remotefile:localpath],example: download:/root/a.txt|/root`)
	flag.IntVar(&identify, "i", 1, `identify: Upload or Download File IdentifyMode
0: direct coverage
1: Overwrite if the file size is different or the modification time is newer, otherwise skip it
2: If the remote file already exists, ask whether to overwrite
`)
	flag.Usage = func() {
		fmt.Println("Version: 0.0.1 BuildTime:2022-4-21 15:35:37")
		flag.PrintDefaults()
	}
	flag.Parse()

	if len(cmdList) == 0 && len(fileoptlist) == 0 && len(TerminalCmd) == 0 {
		flag.Usage()
		os.Exit(1)
	}

	if password == "" {
		password = "xxxxxx"
	}
	cli := Cli{IP: host, Port: port, Username: user, Password: password}
	if len(cmdList) > 0 {
		for _, cmd := range cmdList {
			fmt.Println("command:", cmd)
			output, err := cli.Run(cmd)
			if err != nil {
				fmt.Println(err)
			}
			fmt.Printf("%v", output)
		}
	}

	if TerminalCmd != "" {
		err := cli.RunTerminal(TerminalCmd, os.Stdin, os.Stdout, os.Stderr)
		if err != nil {
			log.Fatal(err)
		}
	}

	if len(fileoptlist) > 0 {
		for _, fileopt := range fileoptlist {
			if strings.HasPrefix(fileopt, "upload:") {
				lt_fname := strings.Split(fileopt[7:], "|")
				if len(lt_fname) == 2 {
					log.Println("---Files upload:", fileopt[7:])
					err := cli.UploadFile(lt_fname[0], lt_fname[1], identify)
					if err != nil {
						log.Println("---Files upload:", fileopt[7:], "failed,", err)
					} else {
						log.Println("---Files upload:", fileopt[7:], "complete")
					}
				} else {
					log.Printf("---Files upload: operation params error,%v\n", fileopt)
				}
			} else if strings.HasPrefix(fileopt, "download:") {
				log.Println("---Files download:", fileopt[9:])
				lt_fname := strings.Split(fileopt[9:], "|")
				if len(lt_fname) == 2 {
					err := cli.DownloadFile(lt_fname[1], lt_fname[0], identify)
					if err != nil {
						log.Println("---Files download:", fileopt[9:], "failed,", err)
					} else {
						log.Println("---Files download:", fileopt[9:], " complete")
					}
				} else {
					log.Printf("---Files download: operation params error,%v\n", fileopt)
				}
			} else {
				log.Printf("---Files operation: is unsupported,%v\n", fileopt)
			}
		}
	}
	cli.close()
}
